package com.epam.rd.mycollections;

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Objects;

public class MyTreeMap<K, V> implements MyMap<K, V> {

    private static boolean RED = true;
    private static boolean BLACK = false;

    private Comparator<? super K> comparator;
    private SimpleEntry<K, V> root;
    private int size;

    public MyTreeMap() {
    }

    public MyTreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }

    private static final class SimpleEntry<K, V> implements Entry<K, V> {
        private K key;
        private V value;
        private SimpleEntry<K, V> left;
        private SimpleEntry<K, V> right;
        private SimpleEntry<K, V> parent;
        private boolean color;

        SimpleEntry(K key, V value, SimpleEntry<K, V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        @Override
        public V getValue() {
            return value;
        }

        @Override
        public K getKey() {
            return key;
        }

        @Override
        public V setValue(V newValue) {
            V deletedValue = this.value;
            this.value = newValue;
            return deletedValue;
        }

    }

    @Override
    public void clear() {
        root = null;
        size = 0;
    }

    @Override
    public boolean containsKey(K key) {
        return getEntryFromTree(key) != null;
    }

    private SimpleEntry<K, V> getEntryFromTree(K key) {
        Comparable<? super K> keyToCompare = null;
        if (comparator == null) {
            keyToCompare = (Comparable<? super K>) key;
        }
        SimpleEntry<K, V> entry = root;
        while (entry != null) {
            int resultOfComparing = compare(keyToCompare, key, entry.key);
            if (resultOfComparing < 0) {
                entry = entry.left;
            } else if (resultOfComparing > 0) {
                entry = entry.right;
            } else {
                return entry;
            }
        }
        return null;
    }

    @Override
    public boolean containsValue(V value) {
        SimpleEntry<K, V> entry = firstEntry();
        do {
            if (entry != null && Objects.equals(value, entry.value)) {
                return true;
            }
            entry = nextAfter(entry);
        } while (entry != null);
        return false;
    }

    private SimpleEntry<K, V> nextAfter(SimpleEntry<K, V> entry) {
        if (entry == null) {
            return null;
        } else if (entry.right != null) {
            SimpleEntry<K, V> next = entry.right;
            while (next.left != null) {
                next = next.left;
            }
            return next;
        } else {
            SimpleEntry<K, V> next = entry.parent;
            SimpleEntry<K, V> pred = entry;
            while (next != null && pred == next.right) {
                pred = next;
                next = next.parent;
            }
            return next;
        }
    }

    private SimpleEntry<K, V> firstEntry() {
        SimpleEntry<K, V> entry = root;
        if (entry != null) {
            while (entry.left != null) {
                entry = entry.left;
            }
        }
        return entry;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public V get(K key) {
        SimpleEntry<K, V> entry = getEntryFromTree(key);
        if (entry != null) {
            return entry.getValue();
        }
        return null;
    }

    @Override
    public V put(K key, V value) {
        if (key == null) {
            throw new NullPointerException();
        }
        SimpleEntry<K, V> entry = root;
        int resultOfCompare;
        SimpleEntry<K, V> parent;
        Comparable<? super K> keyCompare = null;
        SimpleEntry<K, V> addedEntry;
        if (entry == null) {
            root = new SimpleEntry<K, V>(key, value, null);
            size = 1;
            return null;
        }
        if (comparator == null) {
            keyCompare = (Comparable<? super K>) key;
        }
        do {
            parent = entry;
            resultOfCompare = compare(keyCompare, key, entry.key);
            if (resultOfCompare < 0) {
                entry = entry.left;
            } else if (resultOfCompare > 0) {
                entry = entry.right;
            } else {
                return entry.setValue(value);
            }
        } while (entry != null);
        addedEntry = new SimpleEntry<>(key, value, parent);
        if (resultOfCompare < 0) {
            parent.left = addedEntry;
        } else {
            parent.right = addedEntry;
        }
        rebalanceAfterInsert(addedEntry);
        size++;
        return null;
    }

    private int compare(Comparable<? super K> key, K keyOne, K keyTwo) {
        if (comparator == null) {
            return key.compareTo(keyTwo);
        } else {
            return comparator.compare(keyOne, keyTwo);
        }
    }

    private void rebalanceAfterInsert(SimpleEntry<K, V> addedEntry) {
        SimpleEntry<K, V> entry = addedEntry;
        SimpleEntry<K, V> uncle;
        entry.color = RED;
        while (entry != null && entry != root && entry.parent.color == RED) {
            if (parentOf(entry) == leftOf(parentOf(parentOf(entry)))) {
                uncle = rightOf(parentOf(parentOf(entry)));
                if (colorOf(uncle) == BLACK) {
                    if (entry == rightOf(parentOf(entry))) {
                        entry = parentOf(entry);
                        rotateLeft(entry);
                    }
                    setColor(parentOf(entry), BLACK);
                    setColor(parentOf(parentOf(entry)), RED);
                    rotateRight(parentOf(parentOf(entry)));
                } else {
                    entry = setBalance(entry, uncle);
                }
            } else {
                uncle = leftOf(parentOf(parentOf(entry)));
                if (colorOf(uncle) == BLACK) {
                    if (entry == leftOf(parentOf(entry))) {
                        entry = parentOf(entry);
                        rotateRight(entry);
                    }
                    setColor(parentOf(entry), BLACK);
                    setColor(parentOf(parentOf(entry)), RED);
                    rotateLeft(parentOf(parentOf(entry)));
                } else {
                    entry = setBalance(entry, uncle);
                }
            }
        }
        root.color = BLACK;
    }

    private static <K, V> SimpleEntry<K, V> setBalance(SimpleEntry<K, V> entry,
            SimpleEntry<K, V> uncle) {
        setColor(parentOf(entry), BLACK);
        setColor(uncle, BLACK);
        setColor(parentOf(parentOf(entry)), RED);
        return parentOf(parentOf(entry));
    }

    private static <K, V> boolean colorOf(SimpleEntry<K, V> entry) {
        boolean result;
        if (entry == null) {
            result = BLACK;
        } else {
            result = entry.color;
        }
        return result;
    }

    private static <K, V> SimpleEntry<K, V> parentOf(SimpleEntry<K, V> entry) {
        if (entry == null) {
            return null;
        } else {
            return entry.parent;
        }
    }

    private static <K, V> void setColor(
            SimpleEntry<K, V> entry, boolean color) {
        if (entry != null) {
            entry.color = color;
        }
    }

    private static <K, V> SimpleEntry<K, V> leftOf(SimpleEntry<K, V> entry) {
        if (entry == null) {
            return null;
        } else {
            return entry.left;
        }
    }

    private static <K, V> SimpleEntry<K, V> rightOf(SimpleEntry<K, V> entry) {
        if (entry == null) {
            return null;
        } else {
            return entry.right;
        }
    }

    private void rotateLeft(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> right;
        if (entry != null) {
            right = entry.right;
            entry.right = right.left;
            if (right.left != null) {
                right.left.parent = entry;
            }
            right.parent = entry.parent;
            if (entry.parent == null) {
                root = right;
            } else if (entry.parent.left == entry) {
                entry.parent.left = right;
            } else {
                entry.parent.right = right;
            }
            right.left = entry;
            entry.parent = right;
        }
    }

    private void rotateRight(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> left;
        if (entry != null) {
            left = entry.left;
            entry.left = left.right;
            if (left.right != null) {
                left.right.parent = entry;
            }
            left.parent = entry.parent;
            if (entry.parent == null) {
                root = left;
            } else if (entry.parent.right == entry) {
                entry.parent.right = left;
            } else {
                entry.parent.left = left;
            }
            left.right = entry;
            entry.parent = left;
        }
    }

    private void rebalanceAfterDeletion(SimpleEntry<K, V> deletedEntry) {
        SimpleEntry<K, V> entry = deletedEntry;
        SimpleEntry<K, V> neighbour;
        while (entry != root && colorOf(entry) == BLACK) {
            if (entry == leftOf(parentOf(entry))) {
                neighbour = rightOf(parentOf(entry));
                if (colorOf(neighbour) == RED) {
                    setColor(neighbour, BLACK);
                    setColor(parentOf(entry), RED);
                    rotateLeft(parentOf(entry));
                    neighbour = rightOf(parentOf(entry));
                }
                if (colorOf(leftOf(neighbour)) == BLACK
                        && colorOf(rightOf(neighbour)) == BLACK) {
                    setColor(neighbour, RED);
                    entry = parentOf(entry);
                } else {
                    if (colorOf(rightOf(neighbour)) == BLACK) {
                        setColor(leftOf(neighbour), BLACK);
                        setColor(neighbour, RED);
                        rotateRight(neighbour);
                        neighbour = rightOf(parentOf(entry));
                    }
                    setColor(parentOf(entry), BLACK);
                    setColor(neighbour, colorOf(parentOf(entry)));
                    setColor(rightOf(neighbour), BLACK);
                    rotateLeft(parentOf(entry));
                    entry = root;
                }
            } else {
                neighbour = leftOf(parentOf(entry));
                if (colorOf(neighbour) != BLACK) {
                    setColor(parentOf(entry), RED);
                    setColor(neighbour, BLACK);
                    rotateRight(parentOf(entry));
                    neighbour = leftOf(parentOf(entry));
                }
                if (colorOf(rightOf(neighbour)) == BLACK
                        && colorOf(leftOf(neighbour)) == BLACK) {
                    setColor(neighbour, RED);
                    entry = parentOf(entry);
                } else {
                    if (colorOf(leftOf(neighbour)) == BLACK) {
                        setColor(rightOf(neighbour), BLACK);
                        setColor(neighbour, RED);
                        rotateLeft(neighbour);
                        neighbour = leftOf(parentOf(entry));
                    }
                    setColor(leftOf(neighbour), BLACK);
                    setColor(neighbour, colorOf(parentOf(entry)));
                    setColor(parentOf(entry), BLACK);
                    rotateRight(parentOf(entry));
                    entry = root;
                }
            }
        }
        setColor(entry, BLACK);
    }

    private void deleteEntry(SimpleEntry<K, V> deleted) {
        SimpleEntry<K, V> entry = deleted;
        size--;
        if (entry.left != null && entry.right != null) {
            SimpleEntry<K, V> next = nextAfter(entry);
            entry.key = next.getKey();
            entry.value = next.value;
            entry = next;
        }
        SimpleEntry<K, V> replacement;
        if (entry.left != null) {
            replacement = entry.left;
        } else {
            replacement = entry.right;
        }
        if (replacement != null) {
            replacement.parent = entry.parent;
            if (entry.parent == null) {
                root = replacement;
            } else if (entry == entry.parent.left) {
                entry.parent.left = replacement;
            } else {
                entry.parent.right = replacement;
            }
            entry.left = null;
            entry.right = null;
            entry.parent = null;
            if (entry.color == BLACK) {
                rebalanceAfterDeletion(replacement);
            }
        } else if (entry.parent == null) {
            root = null;
        } else {
            if (entry.color == BLACK) {
                rebalanceAfterDeletion(entry);
            }
            if (entry.parent != null) {
                if (entry == entry.parent.left) {
                    entry.parent.left = null;
                } else {
                    entry.parent.right = null;
                }
                entry.parent = null;
            }
        }
    }

    @Override
    public V remove(K key) {
        SimpleEntry<K, V> toDelete = getEntryFromTree(key);
        if (toDelete == null) {
            return null;
        }
        V oldValue = toDelete.value;
        deleteEntry(toDelete);
        return oldValue;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public Iterator<Entry<K, V>> entryIterator() {
        return new EntryIterator();
    }

    private class EntryIterator implements Iterator<Entry<K, V>> {
        private SimpleEntry<K, V> next;
        private SimpleEntry<K, V> lastReturned;

        public EntryIterator() {
            next = firstEntry();
        }

        @Override
        public boolean hasNext() {
            return next != null;
        }

        @Override
        public Entry<K, V> next() {
            SimpleEntry<K, V> entry = next;
            next = nextAfter(entry);
            if (entry == null) {
                throw new NoSuchElementException();
            }
            lastReturned = entry;
            return entry;
        }

        @Override
        public void remove() {
            if (lastReturned == null) {
                throw new IllegalStateException();
            }
            if (lastReturned.left != null && lastReturned.right != null) {
                next = lastReturned;
            }
            deleteEntry(lastReturned);
            lastReturned = null;
        }
    }

}
